# # * custom
# save_dir=configs/custom
# save_path=$save_dir/base.yaml

# python configs/gen_base_config.py \
#     --save_dir $save_dir \
#     --codebook_dim 4 \
#     --codebook_size 16384 \
#     --vae_latent_dim 512 \
#     --ch_mult "1,2,2,4" \
#     --attn_resolutions "" \
#     --size 256 \
#     --disc_start 100000000 \
#     --disc_weight 0 \
#     --bs 8

# python configs/gen_clip_config.py \
#     --base_path $save_path \
#     --clip_pos pre

# python configs/gen_clip_config.py \
#     --base_path $save_path \
#     --clip_pos post

# python configs/gen_inference_config.py \
#     --base_path $save_path


disc_start=100000000 # either max_value meaning no disc ; either 0 meaning with disc 

# * vq-f4
save_dir=configs/vq-f4
save_path=$save_dir/base.yaml

python configs/gen_base_config.py \
    --save_dir $save_dir \
    --codebook_dim 3 \
    --codebook_size 8192 \
    --vae_latent_dim 512 \
    --ch_mult "1,2,4" \
    --attn_resolutions "" \
    --disc_start $disc_start \
    --bs 8

# * vq-f4-noattn
save_dir=configs/vq-f4-noattn
save_path=$save_dir/base.yaml

python configs/gen_base_config.py \
    --save_dir $save_dir \
    --codebook_dim 3 \
    --codebook_size 8192 \
    --vae_latent_dim 512 \
    --ch_mult "1,2,4" \
    --attn_resolutions "" \
    --no_attn \
    --disc_start $disc_start \
    --bs 8


# * vq-f8
save_dir=configs/vq-f8
save_path=$save_dir/base.yaml

python configs/gen_base_config.py \
    --save_dir $save_dir \
    --codebook_dim 4 \
    --codebook_size 16384 \
    --vae_latent_dim 512 \
    --ch_mult "1,2,2,4" \
    --attn_resolutions "32" \
    --disc_num_layers 2 \
    --disc_start $disc_start \
    --disc_weight 0.6 \
    --size 384 \
    --bs 10


# * vq-f8-n256
save_dir=configs/vq-f8-n256
save_path=$save_dir/base.yaml

python configs/gen_base_config.py \
    --save_dir $save_dir \
    --codebook_dim 4 \
    --codebook_size 256 \
    --vae_latent_dim 512 \
    --ch_mult "1,2,2,4" \
    --attn_resolutions "32" \
    --disc_start $disc_start \
    --size 384 \
    --bs 10


# * vq-f16
save_dir=configs/vq-f16
save_path=$save_dir/base.yaml

python configs/gen_base_config.py \
    --save_dir $save_dir \
    --codebook_dim 8 \
    --codebook_size 16384 \
    --vae_latent_dim 512 \
    --ch_mult "1,1,2,2,4" \
    --attn_resolutions "16" \
    --disc_start $disc_start \
    --disc_num_layers 2 \
    --size 384 \
    --bs 14


# * CLIP and inference config
model_types=("vq-f4" "vq-f4-noattn" "vq-f8" "vq-f8-n256" "vq-f16")
for model_type in ${model_types[*]}
do
    save_path=configs/$model_type/base.yaml
    python configs/gen_clip_config.py \
        --base_path $save_path \
        --clip_pos pre

    python configs/gen_clip_config.py \
        --base_path $save_path \
        --clip_pos post

    python configs/gen_inference_config.py \
        --base_path $save_path
done
